#%%
import os
# 查看当前工作目录
print(os.getcwd())

#%%
import torch
import torchvision
from torch.utils.tensorboard import SummaryWriter
from torchvision import datasets, transforms

# 在./runs下写入
# Writer will output to ./runs/ directory by default
writer = SummaryWriter()
# 下载并归一化数据
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
trainset = datasets.MNIST('mnist_train', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)
# 用resnet50当模型
model = torchvision.models.resnet50(False)
#%%
# ResNet需要灰度图像而不是RGB图像
# 所有conv1这一层要改变一下
# Have ResNet model take in grayscale rather than RGB
model.conv1 = torch.nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
# 取一个图像和标签出来看下呢
images, labels = next(iter(trainloader))
# 设置网格把图像装起来
grid = torchvision.utils.make_grid(images)
# 用writer向tensorboard写入图像
writer.add_image('images', grid, 0)
# 写入模型图,然后关闭
writer.add_graph(model, images)
writer.close()

#%%
# 将训练信息写入tensorboard,分类别聚集
from torch.utils.tensorboard import SummaryWriter
import numpy as np

writer = SummaryWriter()

for n_iter in range(100):
    writer.add_scalar('Loss/train', np.random.random(), n_iter)
    writer.add_scalar('Loss/test', np.random.random(), n_iter)
    writer.add_scalar('Accuracy/train', np.random.random(), n_iter)
    writer.add_scalar('Accuracy/test', np.random.random(), n_iter)
#%%
# 我自己也来写一个试试呢
from torch.utils.tensorboard import SummaryWriter
import numpy as np
# 实例化writer
writer = SummaryWriter()
# 编造要写入的训练数据
history = {}
# 编造损失
history['loss'] = {}
history['loss']['train'] = np.linspace(100,1,100)
history['loss']['test'] = np.linspace(90,1,100)
# 编造精度
history['acc'] = {}
history['acc']['train'] = np.linspace(0,90,100)
history['acc']['test'] = np.linspace(0,80,100)
# 按循环写入数据
for i in range(100):
    # 参数的含义:add_scalar,是按数据点写入的.
    # 第一个是标签,第二个是真实的数据,第三个是第几个(数据的index)
    writer.add_scalar('loss/train',history['loss']['train'][i],i)
    writer.add_scalar('loss/test',history['loss']['test'][i],i)
    writer.add_scalar('acc/train',history['acc']['train'][i],i)
    writer.add_scalar('acc/train',history['acc']['test'][i],i)
# cmd使用 tensorboard --logdir=runs看看效果
#%%
# 高级应用,这下我们只需要建立好summarywriter对象就好了,剩下的更加自动化了
from torch.utils.tensorboard import SummaryWriter

# 这个是自动生成的目录,默认是 runs/month_day_year_hour_minute_secend-device-name.
# create a summary writer with automatically generated folder name.
writer = SummaryWriter()
# folder location: runs/May04_22-14-54_s-MacBook-Pro.local/

# 使用指定的特殊名字
# create a summary writer unp.sing the specified folder name.
writer = SummaryWriter("my_experiment")
# folder location: my_experiment

# comment will append at the tail of directoryname
# create a summary writer with comment appended.
writer = SummaryWriter(comment="LR_0.1_BATCH_16")
# folder location: runs/May04_22-14-54_s-MacBook-Pro.localLR_0.1_BATCH_16/
#%%
# 来试一下,改变目录呢
from torch.utils.tensorboard import SummaryWriter
# 我想设置在./log/exp1/ 下面 
writer = SummaryWriter('log/exp1',comment='test')
writer.add_scalar('test',1000000000000,1)
writer.add_scalar('test',1000000000,2)
# 使用tensorboard--logdir=log/exp1 打开
# ok我会使用了
#%%
# 接下来是
# add_scaler
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter('log/exp2')
for i in range(100):
    writer.add_scalar('x*x+x',i*i+i,i)
writer.close()
# 用tensorboard--logdir=log/exp2打开
# %%
# add_scales
# 一次add多个scalars
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter(log_dir='log/exp3')
r = 5
for i in range(100):
    # 第一个参数main_tag,第2个参数a dict of {tag:scalar},第3个参数global_step
    writer.add_scalars('run_14h', {'xnp.sinx':i*np.np.sin(i/r),
                                    'xcosx':i*np.cos(i/r),
                                    'tanx': np.tan(i/r)}, i)
writer.close()
# This call adds three values to the same scalar plot with the tag
# 'run_14h' in TensorBoard's scalar section.
#%%
# 我也要自己写多个scalar
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter(log_dir='log/exp4')
t = np.linspace(1,10,100)
x=10*(2*np.cos(t)-np.cos(2*t))
y=10*(2*np.sin(t)-np.sin(2*t))
for i in range(100):
    writer.add_scalars('heart-shaped-curve',{'x':x[i],'y':y[i]},global_step=i)
writer.close()
# tensorboard--logdir=log